#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 'utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

def writeEcmp_groupRules(p4info_helper, sw, match_fields, action_params):
    """配置ECMP组规则"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_group",
        match_fields={
            "hdr.ipv4.dstAddr": match_fields
        },
        action_name="MyIngress.set_ecmp_select",
        action_params=action_params
        )
    sw.WriteTableEntry(table_entry)
    print(f"  ECMP组规则已安装: {sw.name} -> 目标: {match_fields}, 基数: {action_params['ecmp_base']}, 数量: {action_params['ecmp_count']}")

def writeEcmp_nhopRules(p4info_helper, sw, match_fields, action_params):
    """
    安装ECMP下一跳规则
    """
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_nhop",
        match_fields={
            "meta.ecmp_select": match_fields
        },
        action_name="MyIngress.set_nhop",
        action_params=action_params)
    sw.WriteTableEntry(table_entry)
    print(f"  已安装ECMP下一跳规则: 哈希值{match_fields} -> MAC {action_params['nhop_dmac']}, 端口 {action_params['port']}")

def writeIpv4_lpmRules(p4info_helper, sw, match_fields, action_params):
    """配置IPv4 LPM转发规则"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": match_fields
        },
        action_name="MyIngress.ipv4_forward",
        action_params=action_params
        )
    sw.WriteTableEntry(table_entry)
    print(f"  IPv4 LPM规则已安装: {sw.name} -> 目标: {match_fields}, MAC: {action_params['dstAddr']}, 端口: {action_params['port']}")

def writeSendFrameRules(p4info_helper, sw, egress_port, smac):
    """安装send_frame表规则，用于重写源MAC地址"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.send_frame",
        match_fields={
            "standard_metadata.egress_port": egress_port
        },
        action_name="MyEgress.rewrite_mac",
        action_params={
            "smac": smac
        })
    sw.WriteTableEntry(table_entry)
    print(f"  已安装send_frame规则: {sw.name} -> 端口{egress_port}, MAC {smac}")

def main(p4info_file_path, bmv2_file_path):
    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        print("正在建立与交换机的连接...")
        
        # Create a switch connection object for s1, s2, s3, s4, s5, s6;
        # this is backed by a P4Runtime gRPC connection.
        # Also, dump all P4Runtime messages sent to switch to given txt files.
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')
        s4 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s4',
            address='127.0.0.1:50054',
            device_id=3,
            proto_dump_file='logs/s4-p4runtime-requests.txt')
        s5 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s5',
            address='127.0.0.1:50055',
            device_id=4,
            proto_dump_file='logs/s5-p4runtime-requests.txt')
        s6 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s6',
            address='127.0.0.1:50056',
            device_id=5,
            proto_dump_file='logs/s6-p4runtime-requests.txt')

        print("正在执行主仲裁更新...")
        # Send master arbitration update message to establish this controller as
        # master (required by P4Runtime before performing any other write operation)
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()
        s4.MasterArbitrationUpdate()
        s5.MasterArbitrationUpdate()
        s6.MasterArbitrationUpdate()
        print("主仲裁更新完成")

        print("正在安装P4程序...")
        # Install the P4 program on the switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s4.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s5.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s6.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        print("P4程序安装完成")

        print("\n开始配置交换机规则...")
        
        # The MAC addr "00:00:00:00:0X:0Y" means the port NO.y of switch 'sX'
        # Sending packets from a switch to another switch, we use the MAC addr "00:00:00:0X:0Y:00"
        # When the packet is sent to a host, it needs to have the right destination MAC address.

        print("\n配置s1交换机规则 (ECMP负载均衡):")
        # 配置ECMP组：目标为10.0.2.2的流量使用ECMP
        writeEcmp_groupRules(p4info_helper, s1, ["10.0.2.2", 32], {"ecmp_base": 0, "ecmp_count": 4})
        # 配置4个ECMP下一跳选择
        writeEcmp_nhopRules(p4info_helper, s1, 0, {"nhop_dmac": "00:00:00:00:02:02", "port": 2})
        writeEcmp_nhopRules(p4info_helper, s1, 1, {"nhop_dmac": "00:00:00:00:02:03", "port": 3})
        writeEcmp_nhopRules(p4info_helper, s1, 2, {"nhop_dmac": "00:00:00:00:02:04", "port": 4})
        writeEcmp_nhopRules(p4info_helper, s1, 3, {"nhop_dmac": "00:00:00:00:02:05", "port": 5})
        # 配置回程路由：目标为10.0.1.1的流量直接转发到h1
        writeIpv4_lpmRules(p4info_helper, s1, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        # 配置s1的send_frame规则
        writeSendFrameRules(p4info_helper, s1, 1, "08:00:00:00:01:01")
        writeSendFrameRules(p4info_helper, s1, 2, "08:00:00:00:01:02")
        writeSendFrameRules(p4info_helper, s1, 3, "08:00:00:00:01:03")
        writeSendFrameRules(p4info_helper, s1, 4, "08:00:00:00:01:04")
        writeSendFrameRules(p4info_helper, s1, 5, "08:00:00:00:01:05")

        print("\n配置s6交换机规则:")
        # 到h2的路由
        writeIpv4_lpmRules(p4info_helper, s6, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 1})
        # 回程路由：到h1的流量转发到s1
        writeIpv4_lpmRules(p4info_helper, s6, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 2})
        # 配置s6的send_frame规则
        writeSendFrameRules(p4info_helper, s6, 1, "08:00:00:00:02:06")
        writeSendFrameRules(p4info_helper, s6, 2, "08:00:00:00:02:01")

        print("\n配置中间交换机规则 (s2, s3, s4, s5):")
        # s2, s3, s4, s5规则：简单的IPv4转发
        print("配置s2交换机:")
        writeIpv4_lpmRules(p4info_helper, s2, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s2, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})
        writeSendFrameRules(p4info_helper, s2, 1, "08:00:00:00:02:01")
        writeSendFrameRules(p4info_helper, s2, 2, "08:00:00:00:02:02")

        print("配置s3交换机:")
        writeIpv4_lpmRules(p4info_helper, s3, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s3, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})
        writeSendFrameRules(p4info_helper, s3, 1, "08:00:00:00:03:01")
        writeSendFrameRules(p4info_helper, s3, 2, "08:00:00:00:03:02")

        print("配置s4交换机:")
        writeIpv4_lpmRules(p4info_helper, s4, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s4, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})
        writeSendFrameRules(p4info_helper, s4, 1, "08:00:00:00:04:01")
        writeSendFrameRules(p4info_helper, s4, 2, "08:00:00:00:04:02")

        print("配置s5交换机:")
        writeIpv4_lpmRules(p4info_helper, s5, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s5, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})
        writeSendFrameRules(p4info_helper, s5, 1, "08:00:00:00:05:01")
        writeSendFrameRules(p4info_helper, s5, 2, "08:00:00:00:05:02")

        print("\n✅ 所有交换机规则配置完成!")
        print("ECMP负载均衡已启用:")
        print("  - h1 (10.0.1.1) 到 h2 (10.0.2.2): 通过s1的4条路径进行负载均衡")
        print("  - h2 (10.0.2.2) 到 h1 (10.0.1.1): 通过s6的4条路径进行负载均衡")
        print("  - 中间交换机 (s2, s3, s4, s5): 提供简单IPv4转发")
        
        print("\n控制器运行中... 按Ctrl+C退出")
        while True:
            sleep(1)

    except KeyboardInterrupt:
        print("\n正在关闭控制器...")
    except grpc.RpcError as e:
        print(f"\nP4Runtime gRPC错误: {e}")
        printGrpcError(e)
    except Exception as e:
        print(f"\n意外错误: {e}")
        import traceback
        traceback.print_exc()

    print("正在关闭所有交换机连接...")
    ShutdownAllSwitchConnections()
    print("控制器已关闭")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
